from ultralytics import YOLO
import os
from datetime import datetime
import cv2
import numpy as np

# 加载模型
model = YOLO('models/yolov8n-seg.pt')

# RTSP视频流配置
RTSP_URL = "rtsp://admin:wkkj1234@192.168.200.74:554/Steaming/Channels/1"  # 替换为你的RTSP流地址
VEHICLE_CLASSES = [2, 5, 7]  # 只检测 car=2, bus=5, truck=7
CLASS_NAMES = {2: 'car', 5: 'bus', 7: 'truck'}


def process_rtsp_stream(rtsp_url, output_dir, process_interval=10):
    """
    处理RTSP视频流并进行车辆检测
    Args:
        rtsp_url: RTSP流地址
        output_dir: 输出目录
        process_interval: 处理间隔（每N帧处理一次）
    """
    # 创建视频捕获对象
    cap = cv2.VideoCapture(rtsp_url)

    if not cap.isOpened():
        print(f"错误：无法打开RTSP流 {rtsp_url}")
        return

    print(f"成功连接RTSP流: {rtsp_url}")

    frame_count = 0
    processed_count = 0
    saved_count = 0

    while True:
        ret, frame = cap.read()

        if not ret:
            print("无法读取视频帧，退出...")
            break

        frame_count += 1

        # 每隔process_interval帧处理一次
        if frame_count % process_interval == 0:
            processed_count += 1
            print(f"\n处理第 {frame_count} 帧...")

            try:
                # 使用predict方法进行预测，确保掩码质量
                results = model.predict(
                    source=frame,
                    imgsz=1280,  # 增大推理尺寸提高精度
                    conf=0.5,  # 置信度阈值
                    iou=0.5,  # IOU阈值
                    retina_masks=True,  # 启用高精度掩码
                    stream=True,  # 流模式优化
                    verbose=False,  # 关闭详细输出
                    save=False  # 不自动保存，我们自己处理
                )

                # 处理检测结果
                vehicle_count, visualized_frame = process_and_visualize_detection(results, frame_count, frame)

                # 只有检测到车辆时才保存图片
                if vehicle_count > 0:
                    save_segmented_image(visualized_frame, frame_count, output_dir)
                    saved_count += 1
                    print(f"检测到 {vehicle_count} 辆车辆，已保存图片")
                else:
                    print(f"未检测到车辆，跳过保存")

                # 实时显示结果
                display_results(visualized_frame, frame_count, vehicle_count)

            except Exception as e:
                print(f"处理第 {frame_count} 帧时出错: {e}")
                continue

        # 按'q'退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
    print(f"\n处理完成！共处理 {processed_count} 帧，保存 {saved_count} 张图片")


def process_and_visualize_detection(results, frame_count, original_frame):
    """
    处理检测结果并生成自定义可视化图片
    使用YOLO原生plot方法确保掩码质量
    """
    vehicle_count = 0

    for result in results:
        if result.boxes is not None and len(result.boxes) > 0:
            # 使用YOLO原生的plot方法获取高质量的可视化结果
            plotted_frame = result.plot()  # 这会生成包含所有检测结果的完整图像

            # 统计车辆数量并创建只包含车辆的可视化结果
            visualized_frame = original_frame.copy()

            for i, box in enumerate(result.boxes):
                cls_id = int(box.cls[0])
                conf = box.conf[0].item()
                cls_name = result.names[cls_id]

                # 只处理车辆类别（2,5,7）
                if cls_id in VEHICLE_CLASSES:
                    vehicle_count += 1
                    print(f"  - 车辆检测: {cls_name} (置信度: {conf:.3f})")

                    # 获取掩码数据
                    if result.masks is not None and i < len(result.masks.data):
                        mask = result.masks.data[i].cpu().numpy()

                        # 从plot结果中提取该车辆的掩码区域
                        # 创建一个只包含当前车辆掩码的图像
                        vehicle_mask_frame = np.zeros_like(original_frame)

                        # 应用掩码到vehicle_mask_frame
                        mask_resized = cv2.resize(mask, (original_frame.shape[1], original_frame.shape[0]))
                        mask_bool = mask_resized > 0.3

                        # 从plot结果中提取该车辆的颜色区域
                        vehicle_mask_frame[mask_bool] = plotted_frame[mask_bool]

                        # 将车辆掩码区域合并到最终结果
                        visualized_frame[mask_bool] = vehicle_mask_frame[mask_bool]

                    # 绘制边界框
                    x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
                    color = get_color_for_class(cls_id)
                    cv2.rectangle(visualized_frame, (x1, y1), (x2, y2), color, 2)

                    # 绘制标签
                    label = f"{cls_name} {conf:.2f}"
                    label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                    cv2.rectangle(visualized_frame, (x1, y1 - label_size[1] - 10),
                                  (x1 + label_size[0], y1), color, -1)
                    cv2.putText(visualized_frame, label, (x1, y1 - 5),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

            if vehicle_count > 0:
                print(f"帧 {frame_count}: 检测到 {vehicle_count} 辆车辆")
                return vehicle_count, visualized_frame
            else:
                print(f"帧 {frame_count}: 未检测到车辆")
                return vehicle_count, original_frame
        else:
            print(f"帧 {frame_count}: 未检测到任何目标")
            return 0, original_frame

    return 0, original_frame


def get_color_for_class(cls_id):
    """为不同类别分配不同颜色"""
    color_map = {
        2: (0, 255, 0),  # car - 绿色
        5: (255, 0, 0),  # bus - 蓝色
        7: (0, 0, 255)  # truck - 红色
    }
    return color_map.get(cls_id, (255, 255, 255))


def save_segmented_image(segmented_img, frame_count, output_dir):
    """只有检测到车辆时才保存图片"""
    if segmented_img is not None:
        # 生成带时间戳的文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
        filename = f"vehicle_detected_frame_{frame_count}_{timestamp}.jpg"
        filepath = os.path.join(output_dir, filename)

        # 保存图片（使用高质量设置）
        cv2.imwrite(filepath, segmented_img, [cv2.IMWRITE_JPEG_QUALITY, 95])
        print(f"保存图片: {filename}")


def display_results(segmented_img, frame_count, vehicle_count):
    """实时显示检测结果"""
    if segmented_img is not None:
        # 添加帧计数信息
        display_frame = segmented_img.copy()

        # 根据是否检测到车辆显示不同信息
        if vehicle_count > 0:
            status_text = f'Frame: {frame_count} - Vehicles: {vehicle_count}'
            status_color = (0, 255, 0)  # 绿色
        else:
            status_text = f'Frame: {frame_count} - No Vehicles'
            status_color = (0, 0, 255)  # 红色

        cv2.putText(display_frame, status_text, (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, status_color, 2)

        # 只有检测到车辆时才显示图例
        if vehicle_count > 0:
            legend_y = 70
            cv2.putText(display_frame, 'Legend:', (10, legend_y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
            legend_y += 30

            for cls_id, cls_name in CLASS_NAMES.items():
                color = get_color_for_class(cls_id)
                color_bgr = (int(color[0]), int(color[1]), int(color[2]))
                cv2.putText(display_frame, f'{cls_name}', (10, legend_y),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, color_bgr, 2)
                legend_y += 25

        # 显示结果
        cv2.imshow('Vehicle Detection - Enhanced Mask Coverage', display_frame)


def process_single_image(image_path, output_dir):
    """处理单张图片（使用原生predict方法确保掩码质量）"""
    print("开始预测单张图片...")
    print(f"输入图片: {image_path}")
    print(f"输出目录: {output_dir}")

    # 使用原生predict方法处理单张图片
    results = model.predict(
        source=image_path,
        imgsz=1280,
        conf=0.25,
        iou=0.5,
        retina_masks=True,
        save=False,
        verbose=False
    )

    print("预测完成！")

    vehicle_count = 0

    for result in results:
        if result.boxes is not None:
            # 使用YOLO原生的plot方法
            plotted_frame = result.plot()

            # 统计车辆数量
            for box in result.boxes:
                cls_id = int(box.cls[0])
                if cls_id in VEHICLE_CLASSES:
                    vehicle_count += 1
                    cls_name = result.names[cls_id]
                    conf = box.conf[0].item()
                    print(f"  - 车辆检测: {cls_name} (置信度: {conf:.3f})")

    # 只有检测到车辆时才保存图片
    if vehicle_count > 0:
        # 生成输出文件名
        original_name = os.path.splitext(os.path.basename(image_path))[0]
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"vehicle_detected_{original_name}_{timestamp}.jpg"
        output_path = os.path.join(output_dir, output_filename)

        # 保存YOLO原生的plot结果
        cv2.imwrite(output_path, plotted_frame, [cv2.IMWRITE_JPEG_QUALITY, 95])
        print(f"检测到 {vehicle_count} 辆车辆，图片已保存: {output_filename}")

        # 显示结果
        display_frame = plotted_frame.copy()
        cv2.putText(display_frame, f'Vehicles: {vehicle_count}', (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.imshow('Detection Result - Native YOLO Plot', display_frame)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    else:
        print("未检测到车辆，不保存图片")


def main():
    """主函数"""
    # 创建带时间戳的输出目录
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = f'results/{timestamp}'
    os.makedirs(output_dir, exist_ok=True)

    print("=" * 60)
    print("车辆实例分割检测系统")
    print("=" * 60)
    print(f"目标类别: {[CLASS_NAMES[cls_id] for cls_id in VEHICLE_CLASSES]}")
    print("其他类别的检测框和掩码都不会显示")
    print("只有检测到车辆时才保存图片")
    print("=" * 60)

    print("选择检测模式:")
    print("1. RTSP视频流检测")
    print("2. 单张图片检测")

    choice = input("请输入选择 (1 或 2): ").strip()

    if choice == "1":
        # RTSP视频流检测
        rtsp_url = input("请输入RTSP流地址 (直接回车使用默认地址): ").strip()
        if not rtsp_url:
            rtsp_url = RTSP_URL

        process_interval = input("请输入处理间隔帧数 (默认10): ").strip()
        process_interval = int(process_interval) if process_interval.isdigit() else 10

        print(f"\n开始RTSP视频流检测...")
        print(f"RTSP地址: {rtsp_url}")
        print(f"处理间隔: 每{process_interval}帧处理一次")
        print(f"推理尺寸: 1280 (高精度)")
        print(f"输出目录: {output_dir}")
        print("按 'q' 退出检测")
        print("-" * 40)

        process_rtsp_stream(rtsp_url, output_dir, process_interval)

    elif choice == "2":
        # 单张图片检测
        image_path = input("请输入图片路径 (直接回车使用默认路径): ").strip()
        if not image_path:
            image_path = r'H:\xiaomi\test4\original_rtsp_frame_486_vehicles1_20251119_133922.jpg'

        process_single_image(image_path, output_dir)

    else:
        print("无效选择！")


if __name__ == "__main__":
    main()